power_pooled = function(n,C,alpha,gamma,tau,sigma,pi,f){
  # Critical t-values
  t_alpha = qt(1 - alpha/2,n*C-3) # critical value for Type I  error
  t_gamma = qt(gamma,n*C-3)       # critical value for Type II error
  
  # Initialize distribution statistics
  mu_ind  = rep(0, length(pi))
  eta_ind = rep(0, length(pi))
  
  # Calculate distribution statistics
  for(i in 1:length(pi)){
    mu_ind[i]  = pi[i]   * f[i]
    eta_ind[i] = pi[i]^2 * f[i]
  }
  
  mu   = sum(mu_ind)
  eta  = sum(eta_ind) # eta=E[pi^2] here; eta=eta-mu^2 in the paper.
  psi  = 0; if(pi[1] == 0) psi = f[1]
  muS  = 1 - mu - psi
  etaT = (eta-mu^2)/(1-psi)-(psi/(1-psi)^2)*mu^2
  
  # Variance multipliers
  varN  = tau + sigma
  varCo = (n-1) * tau
  
  ######### Minimum Detectable Effects #########
  # In the case of a non-trivial design
  Var   = 1/(n*C)*(varCo*(1/(psi*(1-psi)) + (1-psi)/(mu ^2)*etaT) + varN*(psi+mu )/(mu *psi))
  VarS  = 1/(n*C)*(varCo*(1/(psi*(1-psi)) + (1-psi)/(muS^2)*etaT) + varN*(psi+muS)/(muS*psi))
  MDE_T = (t_alpha + t_gamma)*Var^0.5
  MDE_S = (t_alpha + t_gamma)*VarS^0.5
  
  # In the case of a pure control
  Var_T = 1/(n*C)*(varCo*(eta-mu^2)/(mu^2*(1-mu)^2) + varN/(mu*(1-mu)))
  MDE_Tonly = (t_alpha+t_gamma)*Var_T^0.5
  
  # In the case of a pure treatment
  Var_S = 1/(n*C)*(varCo*(eta-(1-mu)^2)/(mu^2*(1-mu)^2) + varN/(mu*(1-mu)))
  MDE_Sonly = (t_alpha+t_gamma)*Var_S^0.5
  
  return(list(MDE_T=MDE_T,MDE_S=MDE_S,MDE_Tonly=MDE_Tonly))
}

power_slope = function(n,C,alpha,gamma,tau,sigma,pi,f,j,k){
  # Critical t-values
  t_alpha = qt(1 - alpha/2,n*C-3)  # critical value for Type I  error
  t_gamma = qt(gamma,n*C-3)        # critical value for Type II error
  
  # Initialize distribution statistics
  mu_ind  = rep(0, length(pi))
  p_ind   = rep(0, length(pi))
  
  # Calculate distribution statistics
  for(i in 1:length(pi)){
    mu_ind[i]  = pi[i] * f[i]
    p_ind[i]   = (1-pi[i]) * f[i]
  }
  
  # Variance multipliers
  varN  = tau + sigma
  varCo = (n-1) * tau
  
  ######### Minimum Detectable Effects #########
  Var_T = (varCo*(1/f[j]+1/f[k])+varN*(1/mu_ind[j]+1/mu_ind[k]))/(n*C)
  Var_S = (varCo*(1/f[j]+1/f[k])+varN*(1/p_ind[j]+1/p_ind[k]))/(n*C)
  MDSE_T = ((t_alpha+t_gamma)/(pi[k]-pi[j]))*Var_T^0.5
  MDSE_S = ((t_alpha+t_gamma)/(pi[k]-pi[j]))*Var_S^0.5
  
  return(list(MDSE_T=MDSE_T,MDSE_S=MDSE_S))
}

power_ind = function(n,C,alpha,gamma,tau,sigma,pi,f,p){
  # Critical t-values
  t_alpha = qt(1 - alpha/2,n*C-3)  # critical value for Type I  error
  t_gamma = qt(gamma,n*C-3)        # critical value for Type II error
  
  # Initialize distribution statistics
  mu_ind  = zeros(1,length(pi))
  p_ind   = zeros(1,length(pi))
  
  # Calculate distribution statistics
  for(i in 1:length(pi)){
    mu_ind[i]  = pi[i] * f[i]
    p_ind[i]   = (1-pi[i]) * f[i]
  }
  
  psi  = 0; if(pi[1] == 0) psi = f[1]
  
  # Variance multipliers
  varN  = tau + sigma
  varCo = (n-1)*tau
  varC  = n*tau + sigma
  
  ######### Minimum Detectable Effects #########
  MDE_ind_T = (t_alpha+t_gamma)*(1/(n*C)*(varCo*(1/f[p]+1/psi)+varN*(1/mu_ind[p]+1/psi)))^0.5
  MDE_ind_S = (t_alpha+t_gamma)*(1/(n*C)*(varCo*(1/f[p]+1/psi)+varN*(1/p_ind[p] +1/psi)))^0.5
  return(list(MDE_ind_T=MDE_ind_T,MDE_ind_S=MDE_ind_S))
}

Compass_Search = function(f, admissable, x0, delta = 1, tol=1e-5){
  # INPUTS:
  #    f is a function on domain(x) with list output containing at least $obj to minimize.  Ancillary output is optional.
  #    admissable is a function on domain(x), and returns a boolean if that value is allowed.
  #    x0 is the original guess for the optimal x.
  #    delta is the initial sizes of changes made in guesses for x.
  #    tol is the maximum error tolerated in the final answer for.
  # OUTPUTS:
  #    x is the minimizer of f.
  #    y is the minimum $obj value of f(x).
  #    steps is the number of steps is took to complete the optimization, up to a multiplicative constant.

  # xs = c() # optional object to store working values of x.
  d = length(x0)
  D = rbind(diag(delta,d),diag(-delta,d))
  steps = 0
  x = x0
  while(D[1,1]>tol){
    steps = steps + 1
    admissable_rows = which(apply(D,1,function(row) admissable(x+row))) # check which rows are admissable.
    if(length(admissable_rows) == 0){
      D = D/2
    }else{
      explore_values = apply(matrix(D[admissable_rows,],ncol=d),1,function(row) f(x+row)$obj)
      min_index = order(explore_values)[1]
      no_improvement = explore_values[min_index] >= f(x)$obj # if all values are not an improvement,
      if(no_improvement){
        D = D/2
      }else{
        x = x + D[admissable_rows[min_index],]
      }
    }
    #xs = rbind(xs,x)
  }
  return(list(x=x, y=f(x), steps=steps*d))#,xs=xs
}

# Compass_Search(f=function(x) list(obj=sum(x**2)), admissable=function(x){T}, x0=rep(2.777,5))

index_min = function(x){
  # index of minimum of matrix x.  This can be used for brute search, which is commented out here.
  i = 1 # row
  j = 1 # column
  rows    = nrow(x)
  columns = ncol(x)
  for(r in 1:rows){
    for(c in 1:columns){
      if(x[r,c] < x[i,j]){
        i = r; j = c
      }
    }
  }
  return(c(i,j))
}




###################################################
#                                                 #
#      Replication Code for Baird et al. 2016     #
#                                                 #
###################################################

###################################################
#                Table 1, Column 1                #
###################################################
# Clustered design
# Half of clusters pure control, half pure treatment, ICC = 0.0

n     = 10         # cluster size
C     = 100        # number of clusters
alpha = 0.05       # significance
gamma = 0.8        # power
tau   = 0          # variance of cluster error (tau^2)
sigma = 1          # variance of individual error (sigma^2)
pi    = c(0,1)     # cluster saturations
f     = c(1/2,1/2) # fraction of clusters in each saturation, respectively

X = power_pooled(n,C,alpha,gamma,tau,sigma,pi,f)
X$MDE_Tonly


###################################################
#                Table 1, Column 2                #
###################################################
# Clustered design
# Half of clusters pure control, half pure treatment, ICC = 0.1

n     = 10         # cluster size
C     = 100        # number of clusters
alpha = 0.05       # significance
gamma = 0.8        # power
tau   = 0.1        # variance of cluster error (tau^2)
sigma = 0.9        # variance of individual error (sigma^2)
pi    = c(0,1)     # cluster saturations
f     = c(1/2,1/2) # fraction of clusters in each saturation, respectively

X = power_pooled(n,C,alpha,gamma,tau,sigma,pi,f)
X$MDE_Tonly


###################################################
#                Table 1, Column 3                #
###################################################
# Clustered design
# Half of clusters pure control, half pure treatment, ICC = 1.0


n     = 10         # cluster size
C     = 100        # number of clusters
alpha = 0.05       # significance
gamma = 0.8        # power
tau   = 1.0        # variance of cluster error (tau^2)
sigma = 0.0        # variance of individual error (sigma^2)
pi    = c(0,1)     # cluster saturations
f     = c(1/2,1/2) # fraction of clusters in each saturation, respectively

X = power_pooled(n,C,alpha,gamma,tau,sigma,pi,f)
X$MDE_Tonly


###################################################
#                Table 1, Column 4                #
###################################################
# Partial Population design
# Half of clusters pure control, half partially treated, ICC = 0.0

n     = 10        # cluster size
C     = 100       # number of clusters
alpha = 0.05      # significance
gamma = 0.8       # power
tau   = 0         # variance of cluster error (tau^2)
sigma = 1         # variance of individual error (sigma^2)
pi    = c(0,0.5)  # cluster saturations
theta = 1/2       # proportion of optimization weight on MDE_T

# BRUTE FORCE SEARCH
# psi   = matrix(seq(0.01,.99,by=.001))
# X = matrix(NA,length(psi),3)
# colnames(X) = c("MDE_T", "MDE_S", "MDE_Tonly")
# 
# for(j in 1:length(psi)){
#   f = c(psi[j],1-psi[j])
#   X[j,] = unlist(power_pooled(n,C,alpha,gamma,tau,sigma,pi,f))
# }
# 
# obj   = matrix(theta*X[,"MDE_T"] + (1-theta)*X[,"MDE_S"]) # objective function to maximize
# x = index_min(obj) # index of psi that minimizes objective
# psi[x][1]             # optimal proportion of pure control clusters
# 1 - psi[x][1]         # optimal proportion of partial treatment clusters
# X[x[1],"MDE_T"]       # MDSE_T at optimal control size
# X[x[1],"MDE_S"]       # MDSE_S at optimal control size

# COMPASS SEARCH
admissable = function(x) 0<x[1] & x[1]<1
objective = function(x){# x contains the first saturation level.
  f = c(x[1],1-x[1])
  value = power_pooled(n,C,alpha,gamma,tau,sigma,pi,f)
  objective_function = theta*value$MDE_T + (1-theta)*value$MDE_S # objective function to maximize: a linear combination of MDE_T and MDE_S.
  return(list(obj=objective_function,anc=list(optMDE=value)))
}
opt = Compass_Search(objective, admissable, x0 = c(.5))
opt$x[1]     # optimal proportion of pure control clusters
1-opt$x[1]   # optimal proportion of partial treatment clusters
opt$y$anc$optMDE$MDE_T # MDE_T at optimal control size
opt$y$anc$optMDE$MDE_S # MDE_S at optimal control size

###################################################
#                Table 1, Column 5                #
###################################################
# Partial Population design
# Half of clusters pure control, half partially treated, ICC = 0.1

n     = 10        # cluster size
C     = 100       # number of clusters
alpha = 0.05      # significance
gamma = 0.8       # power
tau   = 0.1       # variance of cluster error (tau^2)
sigma = 0.9       # variance of individual error (sigma^2)
pi    = c(0,0.5)  # cluster saturations
theta = 1/2       # proportion of optimization weight on MDE_T

# BRUTE FORCE SEARCH
# psi   = matrix(seq(0.01,.99,by=.001))
# X = matrix(NA,length(psi),3)
# colnames(X) = c("MDE_T", "MDE_S", "MDE_Tonly")
# for(j in 1:length(psi)){
#   f = c(psi[j],1-psi[j])
#   X[j,] = unlist(power_pooled(n,C,alpha,gamma,tau,sigma,pi,f))
# }

# obj   = matrix(theta*X[,"MDE_T"] + (1-theta)*X[,"MDE_S"]) # objective function to maximize
# x = index_min(obj)    # index of psi that minimizes objective
# psi[x][1]             # optimal proportion of pure control clusters
# 1 - psi[x][1]         # optimal proportion of partial treatment clusters
# X[x[1],"MDE_T"]       # MDE_T at optimal control size
# X[x[1],"MDE_S"]       # MDE_S at optimal control size

# COMPASS SEARCH
admissable = function(x) 0<x[1] & x[1]<1
objective = function(x){# x contains the first saturation level.
  f = c(x[1],1-x[1])
  value = power_pooled(n,C,alpha,gamma,tau,sigma,pi,f)
  objective_function = theta*value$MDE_T + (1-theta)*value$MDE_S # objective function to maximize: a linear combination of MDE_T and MDE_S.
  return(list(obj=objective_function,anc=list(optMDE=value)))
}
opt = Compass_Search(objective, admissable, x0 = c(.5))
opt$x[1]     # optimal proportion of pure control clusters
1-opt$x[1]   # optimal proportion of partial treatment clusters
opt$y$anc$optMDE$MDE_T # MDE_T at optimal control size
opt$y$anc$optMDE$MDE_S # MDE_S at optimal control size



###################################################
#                Table 1, Column 6                #
###################################################
# Partial Population design
# Half of clusters pure control, half partially treated, ICC = 1.0

n     = 10        # cluster size
C     = 100       # number of clusters
alpha = 0.05      # significance
gamma = 0.8       # power
tau   = 1.0       # variance of cluster error (tau^2)
sigma = 0.0       # variance of individual error (sigma^2)
pi    = c(0,0.5)  # cluster saturations
theta = 1/2       # proportion of optimization weight on MDE_T

# BRUTE FORCE SEARCH
# psi   = matrix(seq(0.01,.99,by=.001))
# X = matrix(NA,length(psi),3)
# colnames(X) = c("MDE_T", "MDE_S", "MDE_Tonly")
# for(j in 1:length(psi)){
#   f = c(psi[j],1-psi[j])
#   X[j,] = unlist(power_pooled(n,C,alpha,gamma,tau,sigma,pi,f))
# }
# 
# obj   = matrix(theta*X[,"MDE_T"] + (1-theta)*X[,"MDE_S"]) # objective function to maximize
# ind = index_min(obj)    # index of psi that minimizes objective
# psi[ind][1]             # optimal proportion of pure control clusters
# 1 - psi[ind][1]         # optimal proportion of partial treatment clusters
# X[ind[1],"MDE_T"]       # MDE_T at optimal control size
# X[ind[1],"MDE_S"]       # MDE_S at optimal control size

# COMPASS SEARCH
admissable = function(x) 0<x[1] & x[1]<1
objective = function(x){# x contains the first saturation level.
  f = c(x[1],1-x[1])
  value = power_pooled(n,C,alpha,gamma,tau,sigma,pi,f)
  objective_function = theta*value$MDE_T + (1-theta)*value$MDE_S # objective function to maximize: a linear combination of MDE_T and MDE_S.
  return(list(obj=objective_function,anc=list(optMDE=value)))
}
opt = Compass_Search(objective, admissable, x0 = c(.5))
opt$x[1]     # optimal proportion of pure control clusters
1-opt$x[1]   # optimal proportion of partial treatment clusters
opt$y$anc$optMDE$MDE_T # MDE_T at optimal control size
opt$y$anc$optMDE$MDE_S # MDE_S at optimal control size


###################################################
#                Table 1, Column 7-8              #
###################################################
#### Table 1, Column 7
# Partial Population design
# Some clusters pure control, the rest partially treated, ICC = 0.1
# Subject to twice as much weight on MDE_S compared to MDE_T

n     = 10        # cluster size
C     = 100       # number of clusters
alpha = 0.05      # significance
gamma = 0.8       # power
tau   = 0.1       # variance of cluster error (tau^2)
sigma = 0.9       # variance of individual error (sigma^2)
theta = 1/3       # proportion of optimization weight on MDE_T.

# BRUTE FORCE SEARCH
# x     = matrix(seq(0.01,.99,by=.005)) # values for optimization
# MDE_T = matrix(0,length(x),length(x))
# MDE_S = MDE_T
# for(i in 1:length(x)){
#   pi = c(0,x[i])
#   for (j in 1:length(x)){
#     f = c(x[j],1-x[j])
#     X = power_pooled(n,C,alpha,gamma,tau,sigma,pi,f)
#     MDE_T[i,j] = X$MDE_T
#     MDE_S[i,j] = X$MDE_S
#   }
# }
#   

# obj   = theta*MDE_T + (1-theta)*MDE_S # objective function to maximize
# ind = index_min(obj)  # index of psi that minimizes MDE_T
# x[ind[1]]             # optimal saturation of treatment
# x[ind[2]]             # optimal proportion of pure control clusters
# 1-x[ind[2]]           # optimal proportion of partial treatment clusters
# MDE_T[ind[1],ind[2]]  # MDE_T at optimal control size
# MDE_S[ind[1],ind[2]]  # MDE_T at optimal control size

# COMPASS SEARCH
admissable = function(x) 0<x[1] & x[1]<1 & 0<x[2] & x[2]<1
objective = function(x){# x contains the second saturation level, and the number of people in the pure control group.
  pi = c(0,x[1])
  f = c(x[2],1-x[2])
  value = power_pooled(n,C,alpha,gamma,tau,sigma,pi,f)
  objective_function = theta*value$MDE_T + (1-theta)*value$MDE_S # objective function to maximize: a linear combination of MDE_T and MDE_S.
  return(list(obj=objective_function,anc=list(optMDE=value)))
}
opt = Compass_Search(objective, admissable, x0 = c(.5,.5))
opt$x[1]     # optimal saturation of treatment
opt$x[2]     # optimal proportion of pure control clusters
1-opt$x[2]   # optimal proportion of partial treatment clusters
opt$y$anc$optMDE$MDE_T # MDE_T at optimal control size
opt$y$anc$optMDE$MDE_S # MDE_T at optimal control size



#### Table 1, Column 8
# Partial Population design
# Some clusters pure control, the rest partially treated, ICC = 0.1
# Subject to MDE_T <= 0.25

# BRUTE FORCE SEARCH
# MDE_S_aug = MDE_S;
# for(i in 1:length(x)){
#   for(j in 1:length(x)){
#     if(MDE_T[i,j] > 0.25){
#       MDE_S_aug[i,j] = Inf
#     }
#   }
# }
# obj = MDE_S_aug;      # objective function to maximize
# ind = index_min(obj)  # index of psi that minimizes MDE_S
# x[ind[1]]             # optimal saturation of treatment
# x[ind[2]]             # optimal proportion of pure control clusters
# 1 - x[ind[2]]         # optimal proportion of partial treatment clusters
# MDE_T[ind[1], ind[2]] # MDE_T at optimal control size
# MDE_S[ind[1], ind[2]] # MDE_T at optimal control size

# COMPASS SEARCH
minimum_admissable = function(x) 0<x & x<1
minimum_objective = function(theta){
  admissable = function(x) 0<x[1] & x[1]<1 & 0<x[2] & x[2]<1
  objective = function(x){# x contains the second saturation level, and the number of people in the pure control group.
    pi = c(0,x[1])
    f = c(x[2],1-x[2])
    pooled = power_pooled(n,C,alpha,gamma,tau,sigma,pi,f)
    objective_function = theta*pooled$MDE_T + (1-theta)*pooled$MDE_S
    return(list(obj=objective_function,anc=list(x=x,pooled=pooled)))
  }
  opt = Compass_Search(objective, admissable, x0 = c(0.5,  0.5))
  return(list(obj=abs(opt$y$anc$pooled$MDE_T-0.25),anc=list(opt=opt)))
}
theta = Compass_Search(minimum_objective, minimum_admissable, x0=c(.5))
theta$y$anc$opt$x[1]   # optimal saturation of treatment
theta$y$anc$opt$x[2]   # optimal proportion of pure control clusters
1-theta$y$anc$opt$x[2] # optimal proportion of partial treatment clusters
theta$y$anc$opt$y$anc$pooled$MDE_T # MDE_T at optimal control size
theta$y$anc$opt$y$anc$pooled$MDE_S # MDE_S at optimal control size



###################################################
#                Table 2, Column 1                #
###################################################
# Random saturation design (two partially treated clusters), ICC = 0.0
# Optimal cluster saturations and allocation

n      = 10        # cluster size
C      = 100       # number of clusters
alpha  = 0.05      # significance
gamma  = 0.8       # power
tau    = 0.0       # variance of cluster error (tau^2)
sigma  = 1.0       # variance of individual error (sigma^2)
theta  = 1/2       # fraction of weight put onto MDSE_T

# BRUTE FORCE SEARCH
# x      = seq(0.01,.5,by=.01) # values for optimization
# MDSE_T = matrix(0, length(x),length(x))
# MDSE_S = MDSE_T
# 
# for(i in 1:length(x)){
#   pi = c(x[i],1-x[i])
#   for(j in 1:length(x)){
#     f = c(x[j],1-x[j])
#     X = power_slope(n,C,alpha,gamma,tau,sigma,pi,f,1,2)
#     MDSE_T[i,j] = X$MDSE_T
#     MDSE_S[i,j] = X$MDSE_S
#   }
# }
# 
# obj   = theta*MDSE_T + (1-theta)*MDSE_S; # objective function to maximize
# ind = index_min(obj); # index of psi that minimizes objective
# x[ind[1]]         # optimal saturation for lower saturation
# 1 - x[ind[1]]     # optimal saturation for higher saturation
# x[ind[2]]         # optimal fraction of clusters in lower saturation
# 1 - x[ind[2]]     # optimal fraction of clusters in higher saturation
# MDSE_T[ind[1], ind[2]]  # MDSE_T at optimum
# MDSE_S[ind[1], ind[2]]  # MDSE_S at optimum

# COMPASS SEARCH
admissable = function(x) return(0<x[1] & x[1]<.5 & 0<x[2] & x[2]<.5)
objective  = function(x){ # x contains the first saturation level, and the number of people in the least saturated group.
  pi = c(x[1],1-x[1])
  f = c(x[2],1-x[2])
  value = power_slope(n,C,alpha,gamma,tau,sigma,pi,f,1,2)
  objective_function = theta*value$MDSE_T + theta*value$MDSE_S
  return(list(obj=objective_function,anc=list(opt_MDSE=value)))
}
opt = Compass_Search(objective, admissable, x0=c(.25,.25))
opt$x[1]   # optimal saturation for lower saturation
1-opt$x[1] # optimal saturation for higher saturation
opt$x[2]   # optimal fraction of clusters in lower saturation
1-opt$x[2] # optimal fraction of clusters in higher saturation
opt$y$anc$opt_MDSE$MDSE_T # MDSE_T at optimum
opt$y$anc$opt_MDSE$MDSE_S # MDSE_S at optimum





###################################################
#                Table 2, Column 2                #
###################################################
# Random saturation design (two partially treated clusters), ICC = 0.1
# Optimal cluster saturations and allocation

n      = 10        # cluster size
C      = 100       # number of clusters
alpha  = 0.05      # significance
gamma  = 0.8       # power
tau    = 0.1       # variance of cluster error (tau^2)
sigma  = 0.9       # variance of individual error (sigma^2)
theta  = 1/2       # proportion of optimization weight on MDSE_T

# BRUTE FORCE SEARCH
# x      = seq(0.01,.5,by=.01) # values for optimization
# MDSE_T = matrix(0, length(x),length(x))
# MDSE_S = MDSE_T
# for(i in 1:length(x)){
#   pi = c(x[i],1-x[i])
#   for(j in 1:length(x)){
#     f = c(x[j],1-x[j])
#     X = power_slope(n,C,alpha,gamma,tau,sigma,pi,f,1,2)
#     MDSE_T[i,j] = X$MDSE_T
#     MDSE_S[i,j] = X$MDSE_S
#   }
# }

# BRUTE FORCE SEARCH
# obj = theta*MDSE_T + (1-theta)*MDSE_S; # objective function to maximize
# ind = index_min(obj); # index of psi that minimizes objective
# x[ind[1]]         # optimal saturation for lower saturation
# 1 - x[ind[1]]     # optimal saturation for higher saturation
# x[ind[2]]         # optimal fraction of clusters in lower saturation
# 1 - x[ind[2]]     # optimal fraction of clusters in higher saturation
# MDSE_T[ind[1], ind[2]] # MDSE_T at optimum
# MDSE_S[ind[1], ind[2]] # MDSE_S at optimum

# COMPASS SEARCH
admissable = function(x) return(0<x[1] & x[1]<.5 & 0<x[2] & x[2]<.5)
objective  = function(x){ # x contains the first saturation level, and the number of people in the least saturated group.
  pi = c(x[1],1-x[1])
  f = c(x[2],1-x[2])
  value = power_slope(n,C,alpha,gamma,tau,sigma,pi,f,1,2)
  objective_function = theta*value$MDSE_T + theta*value$MDSE_S
  return(list(obj=objective_function,anc=list(opt_MDSE=value)))
}
opt = Compass_Search(objective, admissable, x0=c(.25,.25))

opt$x[1]   # optimal saturation for lower saturation
1-opt$x[1] # optimal saturation for higher saturation
opt$x[2]   # optimal fraction of clusters in lower saturation
1-opt$x[2] # optimal fraction of clusters in higher saturation
opt$y$anc$opt_MDSE$MDSE_T # MDSE_T at optimum
opt$y$anc$opt_MDSE$MDSE_S # MDSE_S at optimum

###################################################
#                Table 2, Column 3                #
###################################################
# Random saturation design (two partially treated clusters), ICC = 1.0
# Optimal cluster saturations and allocation

n      = 10        # cluster size
C      = 100       # number of clusters
alpha  = 0.05      # significance
gamma  = 0.8       # power
tau    = 1.0       # variance of cluster error (tau^2)
sigma  = 0.0       # variance of individual error (sigma^2)
theta = 1/2        # proportion of optimization weight on  MDSE_T

# BRUTE FORCE SEARCH
# x = seq(0.01,.5,by=.01) # values for optimization
# MDSE_T = matrix(0, length(x),length(x))
# MDSE_S = MDSE_T
# for(i in 1:length(x)){
#   pi = c(x[i],1-x[i])
#   for(j in 1:length(x)){
#     f = c(x[j],1-x[j])
#     X = power_slope(n,C,alpha,gamma,tau,sigma,pi,f,1,2)
#     MDSE_T[i,j] = X$MDSE_T
#     MDSE_S[i,j] = X$MDSE_S
#   }
# }
#
# obj   = theta*MDSE_T + (1-theta)*MDSE_S; # objective function to maximize
# ind = index_min(obj); # index of psi that minimizes objective
# x[ind[1]]         # optimal saturation for lower saturation
# 1 - x[ind[1]]     # optimal saturation for higher saturation
# x[ind[2]]         # optimal fraction of clusters in lower saturation
# 1 - x[ind[2]]     # optimal fraction of clusters in higher saturation
# MDSE_T[ind[1], ind[2]] # MDSE_T at optimum 
# MDSE_S[ind[1], ind[2]] # MDSE_S at optimum

# COMPASS SEARCH
admissable = function(x)  0<x[1] & x[1]<.5 & 0<x[2] & x[2]<.5
objective  = function(x){ # x contains the first saturation level, and the number of people in the least saturated group.
  pi = c(x[1],1-x[1])
  f = c(x[2],1-x[2])
  value = power_slope(n,C,alpha,gamma,tau,sigma,pi,f,1,2)
  objective_function = theta*value$MDSE_T + theta*value$MDSE_S
  return(list(obj=objective_function,anc=list(opt_MDSE=value)))
}
opt = Compass_Search(objective, admissable, x0=c(.25,.25))
opt$x[1]   # optimal saturation for lower saturation
1-opt$x[1] # optimal saturation for higher saturation
opt$x[2]   # optimal fraction of clusters in lower saturation
1-opt$x[2] # optimal fraction of clusters in higher saturation
opt$y$anc$opt_MDSE$MDSE_T # MDSE_T at optimum 
opt$y$anc$opt_MDSE$MDSE_S # MDSE_S at optimum 


###################################################
#                Table 2, Column 4                #
###################################################
# Joint maximization across MDE and MDSE.
# NOTE: pi has been maximized using the a joint maximization process,
# performed in Complete_Maximization.m.

n     = 10              # cluster size
C     = 100             # number of clusters
alpha = 0.05            # significance
gamma = 0.8             # power
tau   = 0.1             # variance of cluster error (tau^2)
sigma = 0.9             # variance of individual error (sigma^2)

# BRUTE FORCE SEARCH
#pi    = c(0,0.21,0.88)    # cluster saturations, (given as optimal here).
#x = seq(0.01,.99,by =.01) # values for maximization.
# MDE_T=matrix(1,length(x),length(x))
# MDE_S=MDE_T;MDSE_T=MDE_T;MDSE_S=MDE_T;
# 
# for(i in 1:length(x)){
#   for(j in 1:(length(x)-i)){
#     f=c(x[i],x[j],1-x[i]-x[j]) 
#     X = power_pooled(n,C,alpha,gamma,tau,sigma,pi,f)
#     MDE_T[i,j] = X$MDE_T; MDE_S[i,j] = X$MDE_S;
#     MDSE_T[i,j] = power_slope(n,C,alpha,gamma,tau,sigma,pi,f,2,3)$MDSE_T
#     MDSE_S[i,j] = power_slope(n,C,alpha,gamma,tau,sigma,pi,f,1,3)$MDSE_S
#   }
# }
# 
# obj = MDE_T+MDE_S+MDSE_T+MDSE_S; # objective function to maximize
# ind = index_min(obj);          # indeces of x that minimizes objective function
# c(x[ind[1]], x[ind[2]], 1-x[ind[1]]-x[ind[2]])        # optimal f
# c(MDE_T[ind[1], ind[2]],MDE_S[ind[1], ind[2]],MDSE_T[ind[1], ind[2]],MDSE_S[ind[1], ind[2]]) # optimal MDEs/MDSEs.

# COMPASS SEARCH
admissable = function(x)0<x[1] & x[1] < x[2] & x[2] < 1 & 0<x[3] & 0<x[4] & x[3]+x[4] < 1
objective  = function(x){# x contains the second two saturation levels, and the number of people in the first two control groups.
  pi = c(0,x[1],x[2])
  f = c(x[3],x[4], 1-x[3]-x[4])
  pooled = power_pooled(n,C,alpha,gamma,tau,sigma,pi,f)
  slope1 = power_slope(n,C,alpha,gamma,tau,sigma,pi,f,2,3)$MDSE_T
  slope2 = power_slope(n,C,alpha,gamma,tau,sigma,pi,f,1,3)$MDSE_S
  objective_function = slope1 + slope2 + pooled$MDE_T + pooled$MDE_S
  return(list(obj=objective_function,anc=list(pooled=pooled, slope1=slope1, slope2=slope2)))
}

opt = Compass_Search(objective, admissable, x0=c(.25,.5, .5, .25))
c(0,opt$x[1:2])                 # optimal pi
c(opt$x[3:4],1-sum(opt$x[3:4])) # optimal f
c(opt$y$anc$pooled$MDE_T, opt$y$anc$pooled$MDE_S, opt$y$anc$slope1, opt$y$anc$slope2)

###################################################
#                Table 2, Column 5                #
###################################################
# MDEs resulting from study design of Banerjee et al.

n      = 10        # cluster size
C      = 100       # number of clusters
alpha  = 0.05      # significance
gamma  = 0.8       # power
tau    = 0.1       # variance of cluster error (tau^2) 
sigma  = 0.9       # variance of individual error (sigma^2)
pi     = c(0, 0.25, 0.5, 0.75, 1)   # cluster saturations
f      = c(0.2, 0.2, 0.2, 0.2, 0.2) # fraction of clusters in each saturation
MDE_T  = 0; MDE_S = 0; MDE_Tonly = 0; MDSE_T = 0; MDSE_S = 0
pooled  = power_pooled(n,C,alpha,gamma,tau,sigma,pi,f)
slope_T = power_slope (n,C,alpha,gamma,tau,sigma,pi,f,2,4)
slope_S = power_slope (n,C,alpha,gamma,tau,sigma,pi,f,1,3)
c(pooled$MDE_T, pooled$MDE_S, slope_T$MDSE_T, slope_S$MDSE_S)



###################################################
#                Table 2, Column 6                #
###################################################
# MDEs resulting from study design of Sinclair et al.

n      = 10        # cluster size
C      = 100       # number of clusters
alpha  = 0.05      # significance
gamma  = 0.8       # power
tau    = 0.1       # variance of cluster error (tau^2) 
sigma  = 0.9       # variance of individual error (sigma^2)
pi     = c(0,0.1,0.5,1)         # cluster saturations
f      = c(0.25,0.25,0.25,0.25) # fraction of clusters in each saturation
MDE_T  = 0; MDE_S = 0; MDE_Tonly = 0; MDSE_T = 0; MDSE_S = 0
pooled  = power_pooled(n,C,alpha,gamma,tau,sigma,pi,f)
slope_T = power_slope (n,C,alpha,gamma,tau,sigma,pi,f,2,4)
slope_S = power_slope (n,C,alpha,gamma,tau,sigma,pi,f,1,3)
c(pooled$MDE_T, pooled$MDE_S, slope_T$MDSE_T, slope_S$MDSE_S)



###################################################
#                Table 2, Column 7-8              #
###################################################
#### Table 2, Column 7
# MDEs resulting from study design of Baird et al.

n      = 10        # cluster size
C      = 100       # number of clusters
alpha  = 0.05      # significance
gamma  = 0.8       # power
tau    = 0.1       # variance of cluster error (tau^2) 
sigma  = 0.9       # variance of individual error (sigma^2)
pi    = c(0,0.33,0.67,1)       # cluster saturations
f     = c(0.55,0.15,0.15,0.15) # fraction of clusters in each saturation
MDE_T = 0; MDE_S = 0; MDE_Tonly = 0; MDSE_T = 0; MDSE_S = 0;
pooled  = power_pooled(n,C,alpha,gamma,tau,sigma,pi,f)
slope_T = power_slope (n,C,alpha,gamma,tau,sigma,pi,f,2,4)
slope_S = power_slope (n,C,alpha,gamma,tau,sigma,pi,f,1,3)
c(pooled$MDE_T, pooled$MDE_S, slope_T$MDSE_T, slope_S$MDSE_S)

#### Table 2, Column 8
# Design with optimal MDEs subject to:
# 1.)  MDE_T held no larger than that of Baird et al. (MDE_t <= 0.2679)
# 2.)  Saturation bins held the same as Baird et al.
# 3.)  The share of cluster in two middle saturations held equal.

# BRUTE FORCE SEARCH
# x     = seq(0.01, .99, by=.01) # values for optimization
# MDE_T = matrix(1,length(x),length(x))
# MDE_S = MDE_T; MDSE_T = MDE_T; MDSE_S = MDE_T
# 
# for(i in 1:length(x)){
#   for(j in 1:(length(x)-i)){
#     f = c(x[i],(1-x[i]-x[j])/2,(1-x[i]-x[j])/2, x[j]) 
#     X = power_pooled(n,C,alpha,gamma,tau,sigma,pi,f)
#     MDE_T[i,j] = X$MDE_T; MDE_S[i,j] = X$MDE_S
#     MDSE_T[i,j] = power_slope(n,C,alpha,gamma,tau,sigma,pi,f,2,4)$MDSE_T
#     MDSE_S[i,j] = power_slope(n,C,alpha,gamma,tau,sigma,pi,f,1,3)$MDSE_S
#   }
# }
# 
# MDE_S_aug  = MDE_S
# for(i in 1:length(x)){
#   for(j in 1:length(x)){
#     if(MDE_T[i,j] > 0.2679){
#       MDE_S_aug[i,j] = Inf
#     }
#   }
# }
# obj   = MDE_S_aug          # objective function to maximize
# ind = index_min(obj)       # index of x that minimizes MDE_S
# x[ind[1]]                  # optimal fraction of clusters in pure control
# (1-x[ind[1]]-x[ind[2]])/2  # optimal fraction of clusters in two middle saturations
# x[ind[2]]                  # optimal fraction of clusters in pure treatment
# MDE_T[ind[1], ind[2]]      # MDE_T  at optimal control size
# MDE_S[ind[1], ind[2]]      # MDE_S  at optimal control size
# MDSE_T[ind[1], ind[2]]     # MDSE_T at optimal control size
# MDSE_S[ind[1], ind[2]]     # MDSE_S at optimal control size

# COMPASS SEARCH
minimum_admissable = function(x) 0<x & x<1
minimum_objective  = function(theta){
  admissable = function(x) 0<x[1] & 0<x[2] & x[1]+x[2] < 1
  objective  = function(x){# x contains the first and last saturation level.
    f = c(x[1],(1-x[1]-x[2])/2,(1-x[1]-x[2])/2, x[2]) 
    pooled = power_pooled(n,C,alpha,gamma,tau,sigma,pi,f)
    slope1 = power_slope(n,C,alpha,gamma,tau,sigma,pi,f,2,4)$MDSE_T # MDSE_T at optimal control size
    slope2 = power_slope(n,C,alpha,gamma,tau,sigma,pi,f,1,3)$MDSE_S # MDSE_S at optimal control size
    objective_function = theta*pooled$MDE_T + (1-theta)*pooled$MDE_S
    return(list(obj=objective_function,anc=list(x=x,pooled=pooled, slope1=slope1, slope2=slope2)))
  }
  opt = Compass_Search(objective, admissable, x0 = c(0.25,  0.25))
  return(list(obj=abs(opt$y$anc$pooled$MDE_T-0.2679),anc=list(opt=opt)))
}
theta = Compass_Search(minimum_objective, minimum_admissable, x0=c(1))
theta$y$anc$opt$y$anc$x[1]         # optimal fraction of clusters in pure control
(1-theta$y$anc$opt$y$anc$x[1]-theta$y$anc$opt$y$anc$x[2])/2 # optimal fraction of clusters in two middle saturations
theta$y$anc$opt$y$anc$x[2]         # optimal fraction of clusters in pure treatment
theta$y$anc$opt$y$anc$pooled$MDE_T # MDE_T  at optimal control size
theta$y$anc$opt$y$anc$pooled$MDE_S # MDE_S  at optimal control size
theta$y$anc$opt$y$anc$slope1       # MDSE_T at optimal control size
theta$y$anc$opt$y$anc$slope2       # MDSE_S at optimal control size